import os
import torch
from torch import nn
# from torch.utils import tensorboard

from dataset import train_loader
from model import Model

if __name__ == "__main__":
    model = Model()
    if os.path.exists("output/mnist1.pth"):
        model.load_state_dict(torch.load("output/mnist1.pth"))
    optimizer = torch.optim.Adam(model.parameters())
    # writer = tensorboard.SummaryWriter()

    try:
        for epoch in range(20):
            loss_sum = 0
            for image, label in train_loader:
                optimizer.zero_grad()
                outputs = model(image)
                loss = nn.CrossEntropyLoss()(outputs, label)
                loss.backward()
                optimizer.step()
                loss_sum += loss
            print(loss_sum)
    finally:
        torch.save(model.state_dict(), "./output/mnist2.pth")



    # TEST = True
    # if TEST:
    #     for step, (images, labels) in enumerate(test_loader):
    #         if step == 0:
    #             print(model(images))
